import cv2
import numpy as np
from custom.rectangle import rectangle_process
from matplotlib import pyplot as plt

def ratio_A(cropped_img,thresh):

    # 计算铁片左侧开头位置
    left_edge = np.min(np.where(thresh[thresh.shape[0] // 3, :] == 0))

    # 计算右侧52.1mm的范围
    right_range = int(52.1 / 320 * thresh.shape[1])

    # 计算右侧15.7mm的范围
    right_range_2 = int(15.7 / 320 * thresh.shape[1])

    # 计算右侧区域的像素值0和255的面积
    right_area_0 = np.sum(thresh[:, left_edge + right_range:left_edge + right_range + right_range_2] == 0)
    right_area_255 = np.sum(thresh[:, left_edge + right_range:left_edge + right_range + right_range_2] == 255)

    # 在图像中框出铁片右侧区域
    cv2.rectangle(thresh, (left_edge + right_range, 0), (left_edge + right_range + right_range_2, thresh.shape[0]),
                  (0, 0, 255), 2)

    # 计算面积比值
    ratio = right_area_255 / (right_area_255+right_area_0)

    # 在图像中框出铁片左侧区域
    cv2.rectangle(cropped_img, (left_edge + right_range, 0), (left_edge + right_range + right_range_2, thresh.shape[0]),
                  (0, 0, 255), 2)
    #
    # cv2.imshow('result', cropped_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


    def measure_drop(ratio):
        if ratio <= 0.05:
            return "无脱落"
        elif ratio <= 0.25:
            return "稍有脱落"
        else:
            return "脱落"

    drop = measure_drop(ratio)

    # 输出结果
    print("A区粉尘面积：", right_area_0+right_area_255)
    print("A区脱落面积：", right_area_255)
    print("A区脱落面积占总面积比例：", ratio)
    print("A区脱落程度是:", drop)

    return drop, ratio


def ratio_B(cropped_img,thresh):

    # 计算图像中心位置
    center = np.array(thresh.shape) // 2

    # 计算左侧60.8mm的范围
    left_range = int(60.8 / 320 * thresh.shape[1])

    # 计算左侧区域的像素值0和255的面积
    left_area_0 = np.sum(thresh[:, center[1] - left_range:center[1]] == 0)
    left_area_255 = np.sum(thresh[:, center[1] - left_range:center[1]] == 255)

    # 计算面积比值
    ratio = left_area_255 / (left_area_255+left_area_0)

    # 在图像中框出铁片左侧区域
    cv2.rectangle(cropped_img, (center[1] - left_range, 0), (center[1], thresh.shape[0]), (0, 0, 255), 2)

    # 显示结果图像
    # cv2.imshow('result', cropped_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    def measure_drop(ratio):
        if ratio <= 0.05:
            return "无脱落"
        elif ratio <= 0.25:
            return "稍有脱落"
        else:
            return "脱落"

    drop = measure_drop(ratio)

    # 输出结果
    print("B区粉尘面积：", left_area_0 + left_area_255)
    print("B区脱落面积：", left_area_255)
    print("B区脱落面积占总面积比例：", ratio)
    print("B区脱落程度是:", drop)

    return drop, ratio

def ratio_C(cropped_img,thresh):
    # 计算图像中心位置
    center = np.array(thresh.shape) // 2

    # 计算右侧60.8mm的范围
    right_range = int(60.8 / 320 * thresh.shape[1])

    # 计算右侧47.1mm的范围
    right_range_2 = int(47.1 / 320 * thresh.shape[1])

    # 计算右侧区域的像素值0和255的面积
    right_area_0 = np.sum(thresh[:, center[1] + right_range:center[1] + right_range + right_range_2] == 0)
    right_area_255 = np.sum(thresh[:, center[1] + right_range:center[1] + right_range + right_range_2] == 255)

    # 计算面积比值
    ratio = right_area_255 / (right_area_255 + right_area_0)

    # 在图像中框出铁片右侧区域
    cv2.rectangle(cropped_img, (center[1] + right_range, 0), (center[1] + right_range + right_range_2, thresh.shape[0]),
                  (0, 0, 255), 2)

    # 显示结果图像
    # cv2.imshow('result', cropped_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    def measure_drop(ratio):
        if ratio <= 0.05:
            return "无脱落"
        elif ratio <= 0.25:
            return "稍有脱落"
        else:
            return "脱落"

    drop = measure_drop(ratio)

    # 输出结果
    print("C区粉尘面积：", right_area_0 + right_area_255)
    print("C区脱落面积：", right_area_255)
    print("C区脱落面积占总面积比例：", ratio)
    print("C区脱落程度是:", drop)

    return drop, ratio

def measure_adhesion_level(dropA, dropB, dropC):
        if dropA == '无脱落' and dropB == '无脱落':
            return "A"
        elif dropA == '稍有脱落' and dropB == '无脱落':
            return "B"
        elif dropA == '稍有脱落' and dropB == '稍有脱落':
            return "C"
        elif dropA == '脱落' and dropB == '稍有脱落':
            return "C"
        elif dropB == '脱落' and dropC == '无脱落':
            return "D"
        elif dropB == '脱落' and dropC == '稍有脱落':
            return "E"
        elif dropB == '脱落' and dropC == '脱落':
            return "F"

if __name__ == '__main__':

    img_dir = './imgs/b_xuanzhuan.bmp'
    # img_dir = './imgs/e.bmp'

    img = cv2.imread(img_dir)

    tagged_img, cropped_img = rectangle_process(img)
    # 显示图像
    cv2.imshow('Image', tagged_img)
    # 等待按下任意键
    cv2.waitKey(0)
    # 关闭窗口
    cv2.destroyAllWindows()

    ret, thresh = cv2.threshold(cropped_img, 120, 255, cv2.THRESH_BINARY)

    dropA, ratio_A = ratio_A(cropped_img, thresh)
    dropB, ratio_B = ratio_B(cropped_img, thresh)
    dropC, ratio_C = ratio_C(cropped_img, thresh)

    adhesion_level = measure_adhesion_level(dropA, dropB, dropC)

    # # 输出附着性结果
    print("取向钢涂层附着性等级为：", adhesion_level)
